package com.vova.rec.tensorflow;

import com.google.protobuf.Any;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.tensorflow.framework.*;
import tensorflow.serving.*;

import java.util.*;

@Slf4j
@RestController
@RequestMapping("/test")
public class TestController {
    @Autowired
    private PredictionServiceGrpc.PredictionServiceBlockingStub predictionServiceStub;
    @Autowired
    private ModelServiceGrpc.ModelServiceBlockingStub modelServiceBlockingStub;

    @GetMapping("/status")
    public void status(@RequestParam(defaultValue = "test_model") String modelName) {
        Model.ModelSpec modelSpec = Model.ModelSpec.newBuilder().setName(modelName).build();
        GetModelStatus.GetModelStatusRequest getModelStatusRequest = GetModelStatus.GetModelStatusRequest.newBuilder()
                .setModelSpec(modelSpec).build();
        GetModelStatus.GetModelStatusResponse modelStatus = modelServiceBlockingStub.getModelStatus(getModelStatusRequest);
        log.info("response={}", modelStatus);
    }

    /**
     * 更新模型配置-发布新模型
     */
    @PostMapping("/config")
    public void config() {
        // 设置两个模型配置

        /*
//             日志配置，暂时无效，https://github.com/tensorflow/serving/issues/408
//         INVALID_ARGUMENT: No request-logger-creator provided

        LoggingConfigOuterClass.SamplingConfig samplingConfig = LoggingConfigOuterClass.SamplingConfig.newBuilder()
                .setSamplingRate(0.5).build();  // 请求日志采样率 10%
        LogCollectorConfigOuterClass.LogCollectorConfig logCollectorConfig = LogCollectorConfigOuterClass
                .LogCollectorConfig.newBuilder().setType("info").setFilenamePrefix("./log").build();
        LoggingConfigOuterClass.LoggingConfig loggingConfig = LoggingConfigOuterClass.LoggingConfig.newBuilder()
                .setSamplingConfig(samplingConfig).setLogCollectorConfig(logCollectorConfig).build();
        */

        ModelServerConfigOuterClass.ModelConfig modelConfig1 = ModelServerConfigOuterClass.ModelConfig.newBuilder()
                .setName("test_model").setBasePath("/home/chenjia/tensorflow_server/test_model")
                .setModelPlatform("tensorflow").build();
        ModelServerConfigOuterClass.ModelConfig modelConfig2 = ModelServerConfigOuterClass.ModelConfig.newBuilder()
                .setName("test_model2").setBasePath("/home/chenjia/tensorflow_server/test_model2")
                .setModelPlatform("tensorflow").build();
        ModelServerConfigOuterClass.ModelConfigList modelConfigList = ModelServerConfigOuterClass.ModelConfigList
                .newBuilder().addConfig(modelConfig1).addConfig(modelConfig2).build();
        ModelServerConfigOuterClass.ModelServerConfig modelServerConfig = ModelServerConfigOuterClass.ModelServerConfig
                .newBuilder().setModelConfigList(modelConfigList).build();
        ModelManagement.ReloadConfigRequest reloadConfigRequest = ModelManagement.ReloadConfigRequest.newBuilder()
                .setConfig(modelServerConfig).build();
        ModelManagement.ReloadConfigResponse reloadConfigResponse = modelServiceBlockingStub
                .handleReloadConfigRequest(reloadConfigRequest);
        log.info("reloadConfigResponse={}", reloadConfigResponse);
    }

    @GetMapping("/metadata")
    public void metadata(@RequestParam(defaultValue = "test_model") String modelName) {
        Model.ModelSpec modelSpec = Model.ModelSpec.newBuilder().setName(modelName).build();
        GetModelMetadata.GetModelMetadataRequest getModelMetadataRequest = GetModelMetadata.GetModelMetadataRequest
                .newBuilder().setModelSpec(modelSpec).addMetadataField("signature_def").build();
        GetModelMetadata.GetModelMetadataResponse modelMetadata = predictionServiceStub.getModelMetadata(getModelMetadataRequest);
        log.info("modelMetadata={}", modelMetadata);
        Map<String, Any> metadataMap = modelMetadata.getMetadataMap();
        metadataMap.forEach((key, value) -> {
            log.info("key={},value={}", key, value);
        });
    }

    /**
     * @see <a href="http://docs.api.xiaomi.com/en/cloud-ml/modelservice/use_java_client.html">Java客户端访问</a>
     */
    @GetMapping("/predict")
    public void predict(@RequestParam(defaultValue = "test_model2") String modelName,
                        @RequestParam(defaultValue = "voice_classification") String signatureName,
                        @RequestParam(defaultValue = "dense_input_1:0") String inputSignatureName) {
        /*
        输入 Tensor("dense_input_1:0", shape=(?, 5), dtype=float32)
        输出 Tensor("dense_1_1/Softmax:0", shape=(?, 1), dtype=float32)
         */

        /*
        TensorProto 并不在意输入时的数据格式，只需要在输入时描述好数据类型 DataType 与 TensorShapeProto，然后 server 将会根据数据类型与
            数据形状对入参进行转换，所以，可以自由发挥。
         */
        Model.ModelSpec modelSpec = Model.ModelSpec.newBuilder().setName(modelName).setSignatureName(signatureName).build();
        List<Float> inputs1 = Arrays.asList(
                1.0F, 2.0F, 3.0F, 4.0F, 5.0F,
                11.0F, 3.0F, 5.0F, 10.0F, 52.0F,
                12.0F, 2.2F, 30.0F, 14.0F, 5.0F,
                112.0F, 2.2F, 30.0F, 14.0F, 15.0F);
        TensorShapeProto shapeProto = TensorShapeProto.newBuilder()
                .addDim(TensorShapeProto.Dim.newBuilder().setSize(4).build())
                .addDim(TensorShapeProto.Dim.newBuilder().setSize(5).build()).build();
        TensorProto tensor = TensorProto.newBuilder()
                .setDtype(DataType.DT_FLOAT)
                .setTensorShape(shapeProto)
                .addAllFloatVal(inputs1).build();
        Predict.PredictRequest predictRequest = Predict.PredictRequest.newBuilder().putInputs(inputSignatureName, tensor)
                .setModelSpec(modelSpec).build();
        Predict.PredictResponse predictResponse = predictionServiceStub.predict(predictRequest);
        // 处理预测结果
        log.info("predictResponse={}", predictResponse);
    }

    @Autowired
    private SessionServiceGrpc.SessionServiceBlockingStub sessionServiceBlockingStub;

    public String session() {
        SessionServiceOuterClass.SessionRunRequest sessionRunRequest = SessionServiceOuterClass.SessionRunRequest
                .newBuilder().build();
        sessionServiceBlockingStub.sessionRun(sessionRunRequest);
        return "";
    }


}
